from typing import IO, BinaryIO
import torch
import os


def save_checkpoint(
    model: torch.nn.Module,                     # 要保存的模型
    optimizers: list[torch.optim.Optimizer] | torch.optim.Optimizer,  # 优化器或优化器列表
    iteration: int,                             # 当前训练迭代次数
    out: str | os.PathLike | BinaryIO | IO[bytes],  # 保存路径或文件对象
):
    # 如果模型是经过torch.compile编译的，则提取原始模型
    orig_model = model._orig_mod if hasattr(model, "_orig_mod") else model

    # 确保optimizers是一个列表，如果不是则转换为列表
    if isinstance(optimizers, torch.optim.Optimizer):
        optimizers = [optimizers]

    # 保存检查点，包含模型状态、优化器状态和迭代次数
    torch.save(
        {
            "model": orig_model.state_dict(),  # 保存模型参数
            "optimizer": [optimizer.state_dict() for optimizer in optimizers],  # 保存优化器状态
            "iteration": iteration,            # 保存当前迭代次数
        },
        out,  # 保存到指定路径或文件对象
    )


def load_checkpoint(
    src: str | os.PathLike | BinaryIO | IO[bytes],  # 检查点文件路径或文件对象
    model: torch.nn.Module | None = None,           # 要加载参数的模型（可选）
    optimizers: list[torch.optim.Optimizer] | torch.optim.Optimizer | None = None,  # 要加载状态的优化器或优化器列表（可选）
):
    # 加载检查点文件
    checkpoint = torch.load(src)

    # 如果提供了模型，则加载模型参数
    if model is not None:
        model.load_state_dict(checkpoint["model"])

    # 如果提供了优化器，则加载优化器状态
    if optimizers is not None:
        # 确保optimizers是一个列表
        if isinstance(optimizers, torch.optim.Optimizer):
            optimizers = [optimizers]

        # 遍历优化器列表和检查点中的优化器状态，逐一加载
        for optimizer, state_dict in zip(optimizers, checkpoint["optimizer"]):
            optimizer.load_state_dict(state_dict)

    # 返回保存的迭代次数
    return checkpoint["iteration"]
